{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "AlphaFold_wJackhmmer.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold_wJackhmmer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# AlphaFold2 w/ Jackhmmer (or MMseqs2)\n",
        "\n",
        "---------\n",
        "**UPDATE** (Aug. 13, 2021)\n",
        "\n",
        "This notebook is being retired and no longer updated. The functionality to search using jackhmmer/mmseqs has been integrated in our [new advanced notebook](https://github.com/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced.ipynb).\n",
        "\n",
        "---------\n",
        "\n",
        "This notebook modifies deepmind's [original notebook](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb) to add experimental homooligomer support and option to run MMseqs2 instead of Jackhmmer for MSA generation.\n",
        "\n",
        "See [ColabFold](https://github.com/sokrypton/ColabFold/) for other related notebooks\n",
        "\n",
        "**Limitations**\n",
        "- This notebook does NOT use Templates.\n",
        "- For a typical Google-Colab-GPU (16G) session, the max total length is **1400 residues**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "woIxeCPygt7K",
        "cellView": "form"
      },
      "source": [
        "#@title Install software\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "#@markdown on the left.\n",
        "use_amber_relax = True #@param {type:\"boolean\"}\n",
        "\n",
        "import os\n",
        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '2.0'\n",
        "\n",
        "from IPython.utils import io\n",
        "import subprocess\n",
        "import requests\n",
        "import hashlib\n",
        "import tarfile\n",
        "import time\n",
        "import tqdm.notebook\n",
        "\n",
        "from sys import version_info \n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar'\n",
        "PARAMS_DIR = './alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "# if not already installed\n",
        "if not os.path.isdir(\"alphafold\"):\n",
        "  try:\n",
        "    total = 100 if use_amber_relax else 55\n",
        "    with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      with io.capture_output() as captured:\n",
        "        #######################################################################\n",
        "        %shell rm -rf alphafold\n",
        "        %shell git clone {GIT_REPO} alphafold\n",
        "\n",
        "        # Apply multi-chain patch from Lim Heo @huhlim\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/protein.patch\n",
        "        %shell (patch -u alphafold/alphafold/common/protein.py -i /content/protein.patch)\n",
        "\n",
        "        pbar.update(4)\n",
        "        %shell pip3 install ./alphafold\n",
        "        pbar.update(5)\n",
        "            \n",
        "        %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "        %shell wget -O \"{PARAMS_PATH}\" \"{SOURCE_URL}\"\n",
        "        pbar.update(14)\n",
        "\n",
        "        %shell tar --extract --verbose --file=\"{PARAMS_PATH}\" \\\n",
        "          --directory=\"{PARAMS_DIR}\" --preserve-permissions\n",
        "        %shell rm \"{PARAMS_PATH}\"\n",
        "        pbar.update(27)\n",
        "\n",
        "        #######################################################################\n",
        "        %shell sudo apt install --quiet --yes hmmer\n",
        "        pbar.update(3)\n",
        "\n",
        "        # Install py3dmol.\n",
        "        %shell pip install py3dmol\n",
        "        pbar.update(1)\n",
        "\n",
        "        if use_amber_relax:\n",
        "          # Install OpenMM and pdbfixer.\n",
        "          %shell rm -rf /opt/conda\n",
        "          %shell wget -q -P /tmp \\\n",
        "            https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \\\n",
        "              && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \\\n",
        "              && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "          pbar.update(4)\n",
        "\n",
        "          PATH=%env PATH\n",
        "          %env PATH=/opt/conda/bin:{PATH}\n",
        "          %shell conda update -qy conda \\\n",
        "              && conda install -qy -c conda-forge \\\n",
        "                python={python_version} \\\n",
        "                openmm=7.5.1 \\\n",
        "                pdbfixer\n",
        "          pbar.update(40)\n",
        "\n",
        "          %shell wget -q -P /content \\\n",
        "            https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "          pbar.update(1)\n",
        "          %shell mkdir -p /content/alphafold/common\n",
        "          %shell cp -f /content/stereo_chemical_props.txt /content/alphafold/common\n",
        "\n",
        "          # Apply OpenMM patch.\n",
        "          %shell pushd /opt/conda/lib/python{python_version}/site-packages/ && \\\n",
        "              patch -p0 < /content/alphafold/docker/openmm.patch && \\\n",
        "              popd\n",
        "\n",
        "        # Create a ramdisk to store a database chunk to make Jackhmmer run fast.\n",
        "        %shell sudo mkdir -m 777 --parents /tmp/ramdisk\n",
        "        %shell sudo mount -t tmpfs -o size=9G ramdisk /tmp/ramdisk\n",
        "        pbar.update(1)\n",
        "\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise\n",
        "\n",
        "  ########################################################################################\n",
        "  ########################################################################################\n",
        "\n",
        "import jax\n",
        "if jax.local_devices()[0].platform == 'tpu':\n",
        "  raise RuntimeError('Colab TPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "elif jax.local_devices()[0].platform == 'cpu':\n",
        "  raise RuntimeError('Colab CPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "\n",
        "# --- Python imports ---\n",
        "import sys\n",
        "import pickle\n",
        "if use_amber_relax:\n",
        "  sys.path.append(f\"/opt/conda/lib/python{python_version}/site-packages\")\n",
        "\n",
        "from urllib import request\n",
        "from concurrent import futures\n",
        "from google.colab import files\n",
        "import json\n",
        "from matplotlib import gridspec\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "\n",
        "from alphafold.model import model\n",
        "from alphafold.model import config\n",
        "from alphafold.model import data\n",
        "\n",
        "from alphafold.data import parsers\n",
        "from alphafold.data import pipeline\n",
        "from alphafold.data.tools import jackhmmer\n",
        "\n",
        "from alphafold.common import protein\n",
        "\n",
        "if use_amber_relax:\n",
        "  from alphafold.relax import relax\n",
        "  from alphafold.relax import utils\n",
        "\n",
        "#######################################################################################\n",
        "#######################################################################################\n",
        "def run_mmseqs2(query_sequence, prefix, use_env=True, filter=False):\n",
        "    def submit(query_sequence, mode):\n",
        "      res = requests.post('https://a3m.mmseqs.com/ticket/msa', data={'q':f\">1\\n{query_sequence}\", 'mode': mode})\n",
        "      return res.json()\n",
        "    def status(ID):\n",
        "      res = requests.get(f'https://a3m.mmseqs.com/ticket/{ID}')\n",
        "      return res.json()\n",
        "    def download(ID, path):\n",
        "      res = requests.get(f'https://a3m.mmseqs.com/result/download/{ID}')\n",
        "      with open(path,\"wb\") as out: out.write(res.content)\n",
        "      \n",
        "    if filter:\n",
        "      mode = \"env\" if use_env else \"all\"\n",
        "    else:\n",
        "      mode = \"env-nofilter\" if use_env else \"nofilter\"\n",
        "    \n",
        "    path = f\"{prefix}_{mode}\"\n",
        "    if not os.path.isdir(path): os.mkdir(path)\n",
        "\n",
        "    # call mmseqs2 api\n",
        "    tar_gz_file = f'{path}/out.tar.gz'\n",
        "    if not os.path.isfile(tar_gz_file):\n",
        "      out = submit(query_sequence, mode)\n",
        "      while out[\"status\"] in [\"RUNNING\",\"PENDING\"]:\n",
        "        time.sleep(1)\n",
        "        out = status(out[\"id\"])    \n",
        "      download(out[\"id\"], tar_gz_file)\n",
        "    \n",
        "    # parse a3m files\n",
        "    a3m_lines = []\n",
        "    a3m = f\"{prefix}_{mode}.a3m\"\n",
        "    if not os.path.isfile(a3m):\n",
        "      with tarfile.open(tar_gz_file) as tar_gz: tar_gz.extractall(path)\n",
        "      a3m_files = [f\"{path}/uniref.a3m\"]\n",
        "      if use_env: a3m_files.append(f\"{path}/bfd.mgnify30.metaeuk30.smag30.a3m\")\n",
        "      a3m_out = open(a3m,\"w\")\n",
        "      for a3m_file in a3m_files:\n",
        "        for line in open(a3m_file,\"r\"):\n",
        "          line = line.replace(\"\\x00\",\"\")\n",
        "          if len(line) > 0:\n",
        "            a3m_lines.append(line)\n",
        "            a3m_out.write(line)\n",
        "    else:\n",
        "      a3m_lines = open(a3m).readlines()\n",
        "    return \"\".join(a3m_lines)\n",
        "\n",
        "def run_jackhmmer(sequence, prefix):\n",
        "  pickled_msa_path = f\"{prefix}.jackhmmer.pickle\"\n",
        "  if os.path.isfile(pickled_msa_path):\n",
        "    msas_dict = pickle.load(open(pickled_msa_path,\"rb\"))\n",
        "    msas, deletion_matrices = (msas_dict[k] for k in ['msas', 'deletion_matrices'])\n",
        "    full_msa = []\n",
        "    for msa in msas:\n",
        "      full_msa += msa\n",
        "  else:\n",
        "    # --- Find the closest source ---\n",
        "    test_url_pattern = 'https://storage.googleapis.com/alphafold-colab{:s}/latest/uniref90_2021_03.fasta.1'\n",
        "    ex = futures.ThreadPoolExecutor(3)\n",
        "    def fetch(source):\n",
        "      request.urlretrieve(test_url_pattern.format(source))\n",
        "      return source\n",
        "    fs = [ex.submit(fetch, source) for source in ['', '-europe', '-asia']]\n",
        "    source = None\n",
        "    for f in futures.as_completed(fs):\n",
        "      source = f.result()\n",
        "      ex.shutdown()\n",
        "      break\n",
        "\n",
        "    jackhmmer_binary_path = '/usr/bin/jackhmmer'\n",
        "    dbs = []\n",
        "\n",
        "    num_jackhmmer_chunks = {'uniref90': 59, 'smallbfd': 17, 'mgnify': 71}\n",
        "    total_jackhmmer_chunks = sum(num_jackhmmer_chunks.values())\n",
        "    with tqdm.notebook.tqdm(total=total_jackhmmer_chunks, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      def jackhmmer_chunk_callback(i):\n",
        "        pbar.update(n=1)\n",
        "\n",
        "      pbar.set_description('Searching uniref90')\n",
        "      jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/uniref90_2021_03.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['uniref90'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=135301051)\n",
        "      dbs.append(('uniref90', jackhmmer_uniref90_runner.query('target.fasta')))\n",
        "\n",
        "      pbar.set_description('Searching smallbfd')\n",
        "      jackhmmer_smallbfd_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/bfd-first_non_consensus_sequences.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['smallbfd'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=65984053)\n",
        "      dbs.append(('smallbfd', jackhmmer_smallbfd_runner.query('target.fasta')))\n",
        "\n",
        "      pbar.set_description('Searching mgnify')\n",
        "      jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=jackhmmer_binary_path,\n",
        "          database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/mgy_clusters_2019_05.fasta',\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=num_jackhmmer_chunks['mgnify'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=304820129)\n",
        "      dbs.append(('mgnify', jackhmmer_mgnify_runner.query('target.fasta')))\n",
        "\n",
        "    # --- Extract the MSAs and visualize ---\n",
        "    # Extract the MSAs from the Stockholm files.\n",
        "    # NB: deduplication happens later in pipeline.make_msa_features.\n",
        "\n",
        "    mgnify_max_hits = 501\n",
        "    msas = []\n",
        "    deletion_matrices = []\n",
        "    for db_name, db_results in dbs:\n",
        "      unsorted_results = []\n",
        "      for i, result in enumerate(db_results):\n",
        "        msa, deletion_matrix, target_names = parsers.parse_stockholm(result['sto'])\n",
        "        e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])\n",
        "        e_values = [e_values_dict[t.split('/')[0]] for t in target_names]\n",
        "        zipped_results = zip(msa, deletion_matrix, target_names, e_values)\n",
        "        if i != 0:\n",
        "          # Only take query from the first chunk\n",
        "          zipped_results = [x for x in zipped_results if x[2] != 'query']\n",
        "        unsorted_results.extend(zipped_results)\n",
        "      sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])\n",
        "      db_msas, db_deletion_matrices, _, _ = zip(*sorted_by_evalue)\n",
        "      if db_msas:\n",
        "        if db_name == 'mgnify':\n",
        "          db_msas = db_msas[:mgnify_max_hits]\n",
        "          db_deletion_matrices = db_deletion_matrices[:mgnify_max_hits]\n",
        "        msas.append(db_msas)\n",
        "        deletion_matrices.append(db_deletion_matrices)\n",
        "        msa_size = len(set(db_msas))\n",
        "        print(f'{msa_size} Sequences Found in {db_name}')\n",
        "\n",
        "      pickle.dump({\"msas\":msas,\"deletion_matrices\":deletion_matrices},\n",
        "                  open(pickled_msa_path,\"wb\"))\n",
        "  return msas, deletion_matrices"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W4JpOs6oA-QS"
      },
      "source": [
        "## Making a prediction\n",
        "\n",
        "Please paste the sequence of your protein in the text box below, then run the remaining cells via _Runtime_ > _Run after_. You can also run the cells individually by pressing the _Play_ button on the left.\n",
        "\n",
        "Note that the search against databases and the actual prediction can take some time, from minutes to hours, depending on the length of the protein and what type of GPU you are allocated by Colab (see FAQ below)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "source": [
        "#@title Enter the amino acid sequence to fold ⬇️\n",
        "sequence = 'PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK'  #@param {type:\"string\"}\n",
        "\n",
        "MIN_SEQUENCE_LENGTH = 16\n",
        "MAX_SEQUENCE_LENGTH = 2500\n",
        "\n",
        "# Remove all whitespaces, tabs and end lines; upper-case\n",
        "sequence = sequence.translate(str.maketrans('', '', ' \\n\\t')).upper()\n",
        "#@markdown ### Experimental options\n",
        "homooligomer = 1 #@param [1,2,3,4,5,6,7,8] {type:\"raw\"}\n",
        "\n",
        "full_sequence = sequence * homooligomer\n",
        "\n",
        "aatypes = set('ACDEFGHIKLMNPQRSTVWY')  # 20 standard aatypes\n",
        "if not set(sequence).issubset(aatypes):\n",
        "  raise Exception(f'Input sequence contains non-amino acid letters: {set(sequence) - aatypes}. AlphaFold only supports 20 standard amino acids as inputs.')\n",
        "if len(full_sequence) < MIN_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too short: {len(full_sequence)} amino acids, while the minimum is {MIN_SEQUENCE_LENGTH}')\n",
        "if len(full_sequence) > MAX_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too long: {len(full_sequence)} amino acids, while the maximum is {MAX_SEQUENCE_LENGTH}. Please use the full AlphaFold system for long sequences.')\n",
        "\n",
        "if len(full_sequence) > 1400:\n",
        "  print(f\"WARNING: For a typical Google-Colab-GPU (16G) session, the max total length is 1400 residues. You are at {len(full_sequence)}!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zq5EC6ju-XPo",
        "cellView": "form"
      },
      "source": [
        "#@title Search against genetic databases\n",
        "#@markdown Once this cell has been executed, you will see\n",
        "#@markdown statistics about the multiple sequence alignment \n",
        "#@markdown (MSA) that will be used by AlphaFold. In particular, \n",
        "#@markdown you’ll see how well each residue is covered by similar \n",
        "#@markdown sequences in the MSA.\n",
        "\n",
        "#@markdown ---\n",
        "msa_method = \"jackhmmer\" #@param [\"jackhmmer\",\"mmseqs2\",\"single_sequence\",\"custom_a3m\",\"precomputed\"]\n",
        "#@markdown - `jackhmmer` - default approach from Deepmind\n",
        "#@markdown - `mmseqs2` - fast method from [ColabFold](https://github.com/sokrypton/ColabFold)\n",
        "#@markdown - `single_sequence` - use single sequence input (not recommended, unless a *denovo* design and you dont expect to find any homologous sequences)\n",
        "#@markdown - `custom_a3m` Upload custom MSA (a3m format)\n",
        "#@markdown - `precomputed` If you have previously run this notebook and saved the results,\n",
        "#@markdown you can skip this step by uploading \n",
        "#@markdown the previously generated  `prediction/msa.npz`\n",
        "\n",
        "# prediction directory\n",
        "output_dir = 'prediction'\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "# tmp directory\n",
        "prefix = hashlib.sha1(sequence.encode()).hexdigest()\n",
        "os.makedirs('tmp', exist_ok=True)\n",
        "prefix = os.path.join('tmp',prefix)\n",
        "\n",
        "# --- Search against genetic databases ---\n",
        "with open('target.fasta', 'wt') as f:\n",
        "  f.write(f'>query\\n{sequence}')\n",
        "\n",
        "# Run the search against chunks of genetic databases (since the genetic\n",
        "# databases don't fit in Colab ramdisk).\n",
        "\n",
        "if msa_method == \"precomputed\":\n",
        "  print(\"upload precomputed pickled msa from previous run\")\n",
        "  pickled_msa_dict = files.upload()\n",
        "  msas_dict = pickle.loads(pickled_msa_dict[list(pickled_msa_dict.keys())[0]])\n",
        "  msas, deletion_matrices = (msas_dict[k] for k in ['msas', 'deletion_matrices'])\n",
        "\n",
        "elif msa_method == \"mmseqs2\":\n",
        "  msa, deletion_matrix = parsers.parse_a3m(run_mmseqs2(sequence, prefix, filter=True))\n",
        "  msas,deletion_matrices = [msa],[deletion_matrix]\n",
        "\n",
        "elif msa_method == \"single_sequence\":\n",
        "  msas = [[sequence]]\n",
        "  deletion_matrices = [[[0]*len(sequence)]]\n",
        "\n",
        "elif msa_method == \"custom_a3m\":\n",
        "  print(\"upload custom a3m\")\n",
        "  msa_dict = files.upload()\n",
        "  lines = msa_dict[list(msa_dict.keys())[0]].decode().splitlines()\n",
        "  a3m_lines = []\n",
        "  for line in lines:\n",
        "    line = line.replace(\"\\x00\",\"\")\n",
        "    if len(line) > 0 and not line.startswith('#'):\n",
        "      a3m_lines.append(line)\n",
        "  msa, deletion_matrix = parsers.parse_a3m(\"\\n\".join(a3m_lines))\n",
        "  msas,deletion_matrices = [msa],[deletion_matrix]\n",
        "\n",
        "  if len(msas[0][0]) != len(sequence):\n",
        "    print(\"ERROR: the length of msa does not match input sequence\")\n",
        "\n",
        "else:\n",
        "  # run jackhmmer\n",
        "  msas, deletion_matrices = run_jackhmmer(sequence, prefix)\n",
        "\n",
        "full_msa = []\n",
        "for msa in msas: full_msa += msa\n",
        "\n",
        "# save MSA as pickle\n",
        "pickle.dump({\"msas\":msas,\"deletion_matrices\":deletion_matrices},\n",
        "            open(\"prediction/msa.pickle\",\"wb\"))\n",
        "\n",
        "# deduplicate\n",
        "deduped_full_msa = list(dict.fromkeys(full_msa))\n",
        "total_msa_size = len(deduped_full_msa)\n",
        "print(f'\\n{total_msa_size} Sequences Found in Total\\n')\n",
        "\n",
        "msa_arr = np.array([list(seq) for seq in deduped_full_msa])\n",
        "num_alignments, num_res = msa_arr.shape\n",
        "\n",
        "plt.figure(figsize=(8,5),dpi=100)\n",
        "plt.title(\"Sequence coverage\")\n",
        "seqid = (np.array(list(sequence)) == msa_arr).mean(-1)\n",
        "seqid_sort = seqid.argsort() #[::-1]\n",
        "non_gaps = (msa_arr != \"-\").astype(float)\n",
        "non_gaps[non_gaps == 0] = np.nan\n",
        "plt.imshow(non_gaps[seqid_sort]*seqid[seqid_sort,None],\n",
        "           interpolation='nearest', aspect='auto',\n",
        "           cmap=\"rainbow_r\", vmin=0, vmax=1, origin='lower')\n",
        "plt.plot((msa_arr != \"-\").sum(0), color='black')\n",
        "plt.xlim(-0.5,msa_arr.shape[1]-0.5)\n",
        "plt.ylim(-0.5,msa_arr.shape[0]-0.5)\n",
        "plt.colorbar(label=\"Sequence identity to query\",)\n",
        "plt.xlabel(\"Positions\")\n",
        "plt.ylabel(\"Sequences\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DynHMYHTzXbO",
        "cellView": "form"
      },
      "source": [
        "#@title Run AlphaFold \n",
        "\n",
        "#@markdown ---\n",
        "relax_all = False #@param {type:\"boolean\"}\n",
        "turbo = False #@param {type:\"boolean\"}\n",
        "#@markdown - `relax_all` If disabled, only top ranked model is amber-relaxed.\n",
        "#@markdown - `turbo` mode compiles a single model then swaps out params to speedup calculation. Warning, this option is experimental, and not extensively tested.\n",
        "\n",
        "\n",
        "# --- Run the model ---\n",
        "from string import ascii_uppercase\n",
        "\n",
        "def _placeholder_template_feats(num_templates_, num_res_):\n",
        "  return {\n",
        "      'template_aatype': np.zeros([num_templates_, num_res_, 22], np.float32),\n",
        "      'template_all_atom_masks': np.zeros([num_templates_, num_res_, 37, 3], np.float32),\n",
        "      'template_all_atom_positions': np.zeros([num_templates_, num_res_, 37], np.float32),\n",
        "      'template_domain_names': np.zeros([num_templates_], np.float32),\n",
        "      'template_sum_probs': np.zeros([num_templates_], np.float32),\n",
        "  }\n",
        "\n",
        "num_templates = 0\n",
        "num_res = len(sequence)\n",
        "\n",
        "if homooligomer == 1:\n",
        "  msas_mod = msas\n",
        "  deletion_matrices_mod = deletion_matrices\n",
        "else:\n",
        "  msas_mod = []\n",
        "  deletion_matrices_mod = []\n",
        "  for o in range(homooligomer):\n",
        "    L = num_res * o\n",
        "    R = num_res * (homooligomer-(o+1))\n",
        "    for msa,deletion_matrix in zip(msas,deletion_matrices):\n",
        "      msas_mod.append([\"-\"*L+seq+\"-\"*R for seq in msa])\n",
        "      deletion_matrices_mod.append([[0]*L+mtx+[0]*R for mtx in deletion_matrix])\n",
        "\n",
        "\n",
        "feature_dict = {}\n",
        "feature_dict.update(pipeline.make_sequence_features(sequence*homooligomer, 'test', num_res*homooligomer))\n",
        "feature_dict.update(pipeline.make_msa_features(msas_mod, deletion_matrices=deletion_matrices_mod))\n",
        "feature_dict.update(_placeholder_template_feats(num_templates, num_res*homooligomer))\n",
        "\n",
        "# Minkyung's code\n",
        "# add big enough number to residue index to indicate chain breaks\n",
        "idx_res = feature_dict['residue_index']\n",
        "L_prev = 0\n",
        "Ls = [num_res]*homooligomer\n",
        "for L_i in Ls[:-1]:\n",
        "  idx_res[L_prev+L_i:] += 200\n",
        "  L_prev += L_i  \n",
        "feature_dict['residue_index'] = idx_res\n",
        "\n",
        "plddts = {}\n",
        "pae_outputs = {}\n",
        "unrelaxed_proteins = {}\n",
        "\n",
        "model_names = ['model_4', 'model_1', 'model_2', 'model_3', 'model_5']\n",
        "\n",
        "total = len(model_names)\n",
        "if use_amber_relax:\n",
        "  if relax_all: total += total\n",
        "  else: total += 1\n",
        "\n",
        "with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  for num,model_name in enumerate(model_names):\n",
        "    pbar.set_description(f'Running model_{num+1}')\n",
        "\n",
        "    cfg = config.model_config(model_name+\"_ptm\")\n",
        "    params = data.get_model_haiku_params(model_name+\"_ptm\", './alphafold/data')\n",
        "    if turbo:\n",
        "      #####################################################\n",
        "      # load models or params\n",
        "      if model_name == \"model_4\":\n",
        "        # define model and process features\n",
        "        model_runner = model.RunModel(cfg, params)\n",
        "        processed_feature_dict = model_runner.process_features(feature_dict,random_seed=0)\n",
        "      else:\n",
        "        # swap params\n",
        "        for k in model_runner.params.keys():\n",
        "          model_runner.params[k] = params[k]\n",
        "      prediction_result = model_runner.predict(processed_feature_dict)  \n",
        "      # cleanup to save memory\n",
        "      if model_name == \"model_5\": del model_runner\n",
        "      del params\n",
        "      #####################################################\n",
        "    else:\n",
        "      #####################################################\n",
        "      model_runner = model.RunModel(cfg, params)\n",
        "\n",
        "      processed_feature_dict = model_runner.process_features(feature_dict,random_seed=0)\n",
        "      prediction_result = model_runner.predict(processed_feature_dict)    \n",
        "\n",
        "      # cleanup to save memory\n",
        "      del params\n",
        "      del model_runner\n",
        "      #####################################################\n",
        "\n",
        "    mean_plddt = prediction_result['plddt'].mean()\n",
        "\n",
        "    # Get the pLDDT confidence metrics.\n",
        "    pae_outputs[model_name] = (\n",
        "        prediction_result['predicted_aligned_error'],\n",
        "        prediction_result['max_predicted_aligned_error']\n",
        "    )\n",
        "    plddts[model_name] = prediction_result['plddt']\n",
        "\n",
        "    # Set the b-factors to the per-residue plddt.\n",
        "    final_atom_mask = prediction_result['structure_module']['final_atom_mask']\n",
        "    b_factors = prediction_result['plddt'][:, None] * final_atom_mask\n",
        "    unrelaxed_protein = protein.from_prediction(processed_feature_dict,\n",
        "                                                prediction_result,\n",
        "                                                b_factors=b_factors)\n",
        "    unrelaxed_proteins[model_name] = unrelaxed_protein\n",
        "\n",
        "    # Delete unused outputs to save memory.\n",
        "    del prediction_result\n",
        "    pbar.update(n=1)\n",
        "\n",
        "  # Find the best model according to the mean pLDDT.\n",
        "  model_rank = list(plddts.keys())\n",
        "  model_rank = [model_rank[i] for i in np.argsort([plddts[x].mean() for x in model_rank])[::-1]]\n",
        "\n",
        "  # Write out the prediction\n",
        "  pbar.set_description(f'AMBER relaxation')\n",
        "  for n,name in enumerate(model_rank):\n",
        "    pred_output_path = os.path.join(output_dir,f'model_{n+1}_unrelaxed.pdb')\n",
        "    unrelaxed_protein = unrelaxed_proteins[name]\n",
        "    pdb_lines = protein.to_pdb(unrelaxed_protein)\n",
        "    with open(pred_output_path, 'w') as f:\n",
        "      f.write(pdb_lines)\n",
        "    \n",
        "    if use_amber_relax:\n",
        "      if relax_all or n == 0:\n",
        "        amber_relaxer = relax.AmberRelaxation(\n",
        "            max_iterations=0,\n",
        "            tolerance=2.39,\n",
        "            stiffness=10.0,\n",
        "            exclude_residues=[],\n",
        "            max_outer_iterations=20)\n",
        "        relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=unrelaxed_protein)        \n",
        "        pred_output_path = os.path.join(output_dir,f'model_{n+1}_relaxed.pdb')\n",
        "        with open(pred_output_path, 'w') as f:\n",
        "          f.write(relaxed_pdb_lines)\n",
        "      pbar.update(n=1)\n",
        "      \n",
        "############################################################\n",
        "############################################################\n",
        "\n",
        "for n,rank in enumerate(model_rank):\n",
        "  print(f\"model_{n+1} {plddts[rank].mean()}\")\n",
        "  pae, max_pae = pae_outputs[rank]\n",
        "  # Save pLDDT and predicted aligned error (if it exists)\n",
        "  pae_output_path = os.path.join(output_dir,f'predicted_aligned_error_{n+1}.json')\n",
        "  if pae_outputs:\n",
        "    # Save predicted aligned error in the same format as the AF EMBL DB\n",
        "    rounded_errors = np.round(pae.astype(np.float64), decimals=1)\n",
        "    indices = np.indices((len(rounded_errors), len(rounded_errors))) + 1\n",
        "    indices_1 = indices[0].flatten().tolist()\n",
        "    indices_2 = indices[1].flatten().tolist()\n",
        "    pae_data = json.dumps([{\n",
        "        'residue1': indices_1,\n",
        "        'residue2': indices_2,\n",
        "        'distance': rounded_errors.flatten().tolist(),\n",
        "        'max_predicted_aligned_error': max_pae.item()\n",
        "    }],\n",
        "                          indent=None,\n",
        "                          separators=(',', ':'))\n",
        "    with open(pae_output_path, 'w') as f:\n",
        "      f.write(pae_data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KAZj6CBZTkJM",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "def plot_plddt_legend():\n",
        "  thresh = ['plDDT:','Very low (<50)','Low (60)','OK (70)','Confident (80)','Very high (>90)']\n",
        "  plt.figure(figsize=(1,0.1),dpi=100)\n",
        "  ########################################\n",
        "  for c in [\"#FFFFFF\",\"#FF0000\",\"#FFFF00\",\"#00FF00\",\"#00FFFF\",\"#0000FF\"]:\n",
        "    plt.bar(0, 0, color=c)\n",
        "  plt.legend(thresh, frameon=False,\n",
        "             loc='center', ncol=6,\n",
        "             handletextpad=1,\n",
        "             columnspacing=1,\n",
        "             markerscale=0.5,)\n",
        "  plt.axis(False)\n",
        "  return plt\n",
        "\n",
        "def plot_confidence(model_num=1):\n",
        "  plt.figure(figsize=(10,3),dpi=100)\n",
        "  \"\"\"Plots the legend for plDDT.\"\"\"\n",
        "  #########################################\n",
        "  plt.subplot(1,2,1); plt.title('Predicted lDDT')\n",
        "  plt.plot(plddts[model_rank[model_num-1]])\n",
        "  for n in range(homooligomer+1):\n",
        "    x = n*(len(sequence))\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  #########################################\n",
        "  plt.subplot(1,2,2);plt.title('Predicted Aligned Error')\n",
        "  pae, max_pae = pae_outputs[model_rank[model_num-1]]\n",
        "  plt.imshow(pae, cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "  #########################################\n",
        "  return plt\n",
        "\n",
        "def show_pdb(model_num=1, show_sidechains=False, show_mainchains=False, color=\"lDDT\"):\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  pred_output_path = os.path.join(output_dir,f'model_{model_num}_unrelaxed.pdb')\n",
        "  view.addModel(open(pred_output_path,'r').read(),'pdb')\n",
        "  if color == \"lDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  elif color == \"chain\":\n",
        "    for n,chain,color in zip(range(homooligomer),list(\"ABCDEFGH\"),\n",
        "                     [\"lime\",\"cyan\",\"magenta\",\"yellow\",\"salmon\",\"white\",\"blue\",\"orange\"]):\n",
        "       view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "show_pdb(model_num,show_sidechains, show_mainchains, color).show()\n",
        "if color == \"lDDT\": plot_plddt_legend().show()  \n",
        "plot_confidence(model_num).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Riekgf0KQv_3",
        "cellView": "form"
      },
      "source": [
        "#@title Download prediction\n",
        "\n",
        "#@markdown Once this cell has been executed, a zip-archive with \n",
        "#@markdown the obtained prediction will be automatically downloaded \n",
        "#@markdown to your computer.\n",
        "\n",
        "# add settings file\n",
        "settings_path = os.path.join(output_dir,\"settings.txt\")\n",
        "with open(settings_path, \"w\") as text_file:\n",
        "  text_file.write(f\"sequence={sequence}\\n\")\n",
        "  text_file.write(f\"msa_method={msa_method}\\n\")\n",
        "  text_file.write(f\"homooligomer={homooligomer}\\n\")\n",
        "  text_file.write(f\"use_amber_relax={use_amber_relax}\\n\")\n",
        "  text_file.write(f\"turbo={turbo}\\n\")\n",
        "  text_file.write(f\"use_templates=False\\n\")\n",
        "\n",
        "# --- Download the predictions ---\n",
        "!zip -q -r {output_dir}.zip {output_dir}\n",
        "files.download(f'{output_dir}.zip')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}